package com.shuanwang.nl2sql_demo.nl2Sql.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.shuanwang.nl2sql_demo.nl2Sql.mapper.Nl2SqlMapper;
import com.shuanwang.nl2sql_demo.nl2Sql.service.Nl2SqlService;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Slf4j
@Service
@RequiredArgsConstructor
public class Nl2SqlServiceImpl implements Nl2SqlService {

    private final Nl2SqlMapper nl2SqlMapper;

    private static final String API_KEY = "sk-69f2b0634b304d49abafb1df188d4440";

    @Override
    public List<Map<String, Object>> processQuery(String naturalLanguage) throws Exception {
        // 调用 NL2SQL 服务生成 SQL
        log.info("用户输入的自然语言: 【{}】", naturalLanguage);
        GenerationResult result = callWithMessage(naturalLanguage);

        Optional.ofNullable(result)
                .map(GenerationResult::getOutput)
                .map(GenerationOutput::getChoices)
                .filter(choices -> !CollUtil.isEmpty(choices))
                .orElseThrow(() -> new IllegalArgumentException("无法生成有效的 SQL，请检查输入的自然语言描述。"));

        // 获取生成的 SQL
        String generatedSQL = extractGeneratedSQL(result);

        // 打印生成的 SQL 到日志
        log.info("生成的 SQL: {}", generatedSQL);

        // 使用 MyBatis 执行生成的 SQL 并返回结果
        return executeGeneratedSQL(generatedSQL);
    }

    @Override
    public Object schemaJson() {
        return JSONUtil.parse(loadSchemaFromFile());
    }

    /**
     * 提取生成的 SQL 语句
     */
    private String extractGeneratedSQL(GenerationResult result) {
        return Optional.ofNullable(result)
                .map(GenerationResult::getOutput)
                .map(GenerationOutput::getChoices)
                .filter(choices -> !choices.isEmpty())
                .map(choices -> choices.get(0).getMessage().getContent())
                .filter(StrUtil::isNotBlank)
                .map(String::trim)
                .map(this::extractPureSQL)
                .orElseThrow(() -> new IllegalStateException("生成的 SQL 为空或无效。"));
    }

    /**
     * 提取纯 SQL 部分
     *
     * @param rawContent
     * @return
     */
    private String extractPureSQL(String rawContent) {
        // 正则表达式匹配 SQL 查询
        Pattern sqlPattern = Pattern.compile("(?i)(SELECT.*?;)", Pattern.DOTALL);
        Matcher matcher = sqlPattern.matcher(rawContent);

        if (matcher.find()) {
            return matcher.group(1).trim();
        }
        throw new IllegalArgumentException("未能从生成内容中提取有效的 SQL 查询。");
    }

    /**
     * 执行生成的 SQL 并返回查询结果
     */
    private List<Map<String, Object>> executeGeneratedSQL(String generatedSQL) throws Exception {
        try {
            return nl2SqlMapper.executeQuery(generatedSQL);
        } catch (Exception e) {
            log.error("执行 SQL 时发生错误: {}", generatedSQL, e);
            throw new RuntimeException("SQL 执行失败，请检查生成的 SQL 是否正确。", e);
        }
    }

    private GenerationResult callWithMessage(String generateSQLFromQwen) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你好！我是你的NL2SQL助手，可以帮助你将自然语言查询转换为SQL语句。")
                .build();

        String schemaFromJson = loadSchemaFromFile();

        String prompt = generateSQLFromQwen;

        Message schemaMsgJson = Message.builder()
                .role(Role.USER.getValue())
                .content("表结构说明：" + schemaFromJson + "。请根据问题生成标准的达梦数据库 SQL 查询，不要包含任何解释性文字：")
                .build();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(prompt)
                .build();
        GenerationParam param = GenerationParam.builder()
                .apiKey(API_KEY)
                .model(Generation.Models.QWEN_MAX)
                .messages(Arrays.asList(systemMsg, schemaMsgJson, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        return gen.call(param);
    }

    /**
     * 读取 Schema.json 并返回 JSON 字符串
     *
     * @return
     * @throws Exception
     */
    @SneakyThrows
    private String loadSchemaFromFile() {
        return new String(Files.readAllBytes(Paths.get("SchemaProd.json")));
    }

}